import sys, pdb
from transformers import AutoTokenizer
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
import uvicorn
import get_instruction
from config import Config
config = Config() #获取配置

TIMEOUT_KEEP_ALIVE = 15  # seconds.
app = FastAPI()
engine = None

tokenizer = AutoTokenizer.from_pretrained(config.model_path)

def coach_generate(query, user_id):
    response, _ = get_instruction.coach_response(tokenizer, query, user_id)
    return response

def coach_score_response(user_id):
    global model_url
    response = get_instruction.coach_score_response(tokenizer, user_id)
    return response

'''
对话过程
'''
@app.post("/generate")
@app.get("/generate")
async def req(request: Request):
    #获取请求参数
    if request.method == 'POST':
        request_dict = await request.json()
        query = request_dict.pop('query')
        user_id = request_dict.pop('user_id')
        mode = request_dict.pop('mode')
    else:
        query = request.query_params.get('query')
        user_id = request.query_params.get('user_id')
        mode = request.query_params.get('mode')

    #获取对话
    if mode == 'dialog':
        result = coach_generate(query, user_id)
    elif mode == 'score':
        result = coach_score_response(user_id)
    return result

'''
对话教练
'''
if __name__ == '__main__':
    #python chain.py 13556 http://localhost:12002
    port = int(sys.argv[1])
    model_url = sys.argv[2]
    get_instruction.set_url(model_url)
    uvicorn.run(app,
        host='0.0.0.0',
        port=port,
        log_level="debug",
        timeout_keep_alive=TIMEOUT_KEEP_ALIVE
        )
